import torch.nn as nn
import torch.nn.functional as F
from dgl.nn.pytorch import GraphConv

# 这个类实现了一个简单的逻辑回归模型，用于从隐藏层维度到输出维度的线性映射。
class LogReg(nn.Module):
    def __init__(self, hid_dim, out_dim):
        super(LogReg, self).__init__()
        self.fc = nn.Linear(hid_dim, out_dim)

    def forward(self, x):
        ret = self.fc(x)
        return ret

# 这个类实现了一个多层感知机（MLP），具有两个全连接层和一个可选的批量归一化层。
# use_bn: 是否使用批归一化层的布尔标志，默认值为 True。
class MLP(nn.Module):
    def __init__(self, nfeat, nhid, nclass, use_bn=True):
        super(MLP, self).__init__()

        self.layer1 = nn.Linear(nfeat, nhid, bias=True)
        self.layer2 = nn.Linear(nhid, nclass, bias=True)

        self.bn = nn.BatchNorm1d(nhid)
        self.use_bn = use_bn
        # 定义激活函数，使用 ReLU 激活函数。
        self.act_fn = nn.ReLU()

    def forward(self, _, x):
        x = self.layer1(x)
        if self.use_bn:
            x = self.bn(x)

        x = self.act_fn(x)
        x = self.layer2(x)

        return x

# 这个类实现了一个图卷积网络（GCN），支持多个卷积层。
class GCN(nn.Module):
    def __init__(self, in_dim, hid_dim, out_dim, n_layers):
        super().__init__()
        # 层数
        self.n_layers = n_layers
        # 创建一个ModuleList用于存储每一层的GraphConv
        self.convs = nn.ModuleList()

        # 定义第一层图卷积层，输入维度为in_dim，输出维度为hid_dim
        # norm='both' 表示对输入和输出进行对称归一化。
        # allow_zero_in_degree=True 允许有零入度的节点。
        self.convs.append(GraphConv(in_dim, hid_dim, norm='both', allow_zero_in_degree=True))

        # 如果有多于一层的情况
        if n_layers > 1:
            # 定义中间的图卷积层，输入和输出维度均为hid_dim
            for i in range(n_layers - 2):
                self.convs.append(GraphConv(hid_dim, hid_dim, norm='both', allow_zero_in_degree=True))
            # 定义最后一层图卷积层，输入维度为hid_dim，输出维度为out_dim
            self.convs.append(GraphConv(hid_dim, out_dim, norm='both', allow_zero_in_degree=True))

    def forward(self, graph, attr, x):
      
        for i in range(self.n_layers - 1):
            # n-1层使用ReLU激活函数
            x = F.relu(self.convs[i](graph, x, edge_weight=attr))
        # 最后一层不使用激活函数
        x = self.convs[-1](graph, x, edge_weight=attr)

        return x

# 这个类实现了一个用于图嵌入对比学习的模型，支持使用 GCN 或 MLP 作为骨干网络。
class GCIL(nn.Module):
    def __init__(self, in_dim, hid_dim, out_dim, n_layers, use_mlp = False):
        super().__init__()
        # 判断是否使用 MLP 作为骨干网络
        if not use_mlp:
            # 如果不使用 MLP，则使用 GCN 作为骨干网络
            self.backbone = GCN(in_dim, hid_dim, out_dim, n_layers)
        else:
            self.backbone = MLP(in_dim, hid_dim, out_dim)

    def get_embedding(self, graph, feat, attr):
        # 获取骨干网络的输出并断开梯度计算
        out = self.backbone(graph, attr, feat)
        return out.detach()

    def forward(self, graph1, feat1, attr1, graph2, feat2, attr2):
        # 通过骨干网络计算 graph1 的嵌入
        h1 = self.backbone(graph1, attr1, feat1)
        # 通过骨干网络计算 graph2 的嵌入
        h2 = self.backbone(graph2, attr2, feat2)

        # 对 h1 进行标准化，减去均值除以标准差
        z1 = (h1 - h1.mean(0)) / h1.std(0)
        # 对 h2 进行标准化，减去均值除以标准差
        z2 = (h2 - h2.mean(0)) / h2.std(0)

        # 返回标准化后的嵌入 z1 和 z2 以及原始嵌入 h1 和 h2
        return z1, z2, h1, h2

